#include "ECDSA.h"

#include <ctime>
#include <cassert>

#include "data_conversion.h"


#include "../DEBUG_PRINT_MACRO.h"

#ifdef DEBUG_PRINTS_LEVEL_1
#include <iostream>
#endif

// 引入 Tracy 头文件
#include <Tracy.hpp>

// 辅助函数：计算模逆元，若不可逆则断言失败
void inverse_mod(GMPInt& out, const GMPInt& a, const GMPInt& mod) {
    ZoneScopedN("ECDSA::Helper::inverse_mod"); // 标记辅助函数
    mpz_t inv;
    mpz_init(inv);
    int exists = mpz_invert(inv, a, mod);
    assert(exists != 0 && "Inverse does not exist");
    out = GMPInt(inv);
    mpz_clear(inv);
}


// 判断点是否在曲线上
bool is_on_curve(const EllipticCurve_point& P, const EllipticCurve& curve) {
    ZoneScopedN("ECDSA::Helper::is_on_curve"); // 标记辅助函数
    if (P.is_infinity)
        return true; // 无穷远点始终有效

    GMPInt y_sq;
    {
        ZoneScopedN("ECDSA::Helper::is_on_curve::Calculate Y_sq");
        y_sq = (P.y * P.y) % curve.p;
    }

    GMPInt x_cubed;
    GMPInt ax;
    {
        ZoneScopedN("ECDSA::Helper::is_on_curve::Calculate RHS Components");
        x_cubed = (P.x * P.x * P.x) % curve.p;
        ax = (curve.a * P.x) % curve.p;
    }
    
    GMPInt rhs;
    {
        ZoneScopedN("ECDSA::Helper::is_on_curve::Calculate RHS");
        rhs = (x_cubed + ax + curve.b) % curve.p;
    }
    return (y_sq == rhs);
}

// 点取反
EllipticCurve_point point_neg(const EllipticCurve_point& P, const EllipticCurve& curve) {
    ZoneScopedN("ECDSA::Helper::point_neg"); // 标记辅助函数
    assert(is_on_curve(P, curve)); // 内部已插桩

    if (P.is_infinity)
        return P; // -∞ = ∞
    GMPInt ny = (0 - P.y) % curve.p;
    return EllipticCurve_point(P.x, ny);
}

// 点加法
EllipticCurve_point point_add(const EllipticCurve_point& P, 
                               const EllipticCurve_point& Q,
                               const EllipticCurve& curve) {
    ZoneScopedN("ECDSA::Helper::point_add"); // 标记辅助函数
    assert(is_on_curve(P, curve)); // 内部已插桩
    assert(is_on_curve(Q, curve)); // 内部已插桩

    // 处理无穷远点
    if (P.is_infinity) return Q;
    if (Q.is_infinity) return P;

    // P == -Q 则返回无穷远点
    if (P.x == Q.x && (P.y + Q.y) % curve.p == 0) {
        return EllipticCurve_point::infinity();
    }

    GMPInt m;
    if (P.x == Q.x) { // P == Q，使用切线法
        ZoneScopedN("ECDSA::Helper::point_add::Tangent Method");
        GMPInt numerator = (P.x * P.x * 3 + curve.a) % curve.p;
        GMPInt denominator = (P.y * 2) % curve.p;
        inverse_mod(denominator, denominator, curve.p); // 内部已插桩
        m = (numerator * denominator) % curve.p;
    } else { // P != Q
        ZoneScopedN("ECDSA::Helper::point_add::Secant Method");
        GMPInt dx = (Q.x - P.x) % curve.p;
        GMPInt dy = (Q.y - P.y) % curve.p;
        inverse_mod(dx, dx, curve.p); // 内部已插桩
        m = (dy * dx) % curve.p;
    }

    GMPInt x3 = (m * m - P.x - Q.x) % curve.p;
    GMPInt y3 = (m * (P.x - x3) - P.y) % curve.p;

    return EllipticCurve_point(x3, y3);
}

// 标量乘法（优化双倍加算法）
EllipticCurve_point scalar_mult(const GMPInt& k, 
                                 const EllipticCurve_point& P,
                                 const EllipticCurve& curve) {
    ZoneScopedN("ECDSA::Helper::scalar_mult"); // 标记辅助函数
    assert(is_on_curve(P, curve)); // 内部已插桩

    GMPInt k_mod_n;
    {
        ZoneScopedN("ECDSA::Helper::scalar_mult::Modulo Scalar");
        mpz_mod(k_mod_n, k, curve.n);
    }
    if (k_mod_n == 0 || P.is_infinity)
        return EllipticCurve_point::infinity();

    GMPInt scalar = k;
    if (scalar < 0) {
        scalar = -scalar;
        return scalar_mult(scalar, point_neg(P, curve), curve); // 内部已插桩
    }

    EllipticCurve_point result = EllipticCurve_point::infinity();
    EllipticCurve_point addend = P;

    {
        ZoneScopedN("ECDSA::Helper::scalar_mult::Double-and-Add Loop");
        while (scalar > 0) {
            if (scalar % 2 == 1) {
                result = point_add(result, addend, curve); // 内部已插桩
            }
            addend = point_add(addend, addend, curve); // 内部已插桩
            scalar /= 2;
        }
    }

    assert(is_on_curve(result, curve)); // 内部已插桩
    return result;
}


EllipticCurve get_secp256k1(){
    ZoneScopedN("ECDSA::CurveParams::get_secp256k1"); // 标记曲线参数获取函数
    mpz_t a, b, p, g_x, g_y, n, h;

    {
        ZoneScopedN("ECDSA::CurveParams::get_secp256k1::Initialization");
        // secp256k1曲线参数初始化（实际值预先运算确定）
        mpz_init_set_str(a, "0", 10);        // 曲线参数 a = 0
        mpz_init_set_str(b, "7", 10);        // 曲线参数 b = 7
        mpz_init_set_str(p, "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F", 16); // 素数域p
        mpz_init_set_str(g_x, "79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798", 16); // 基点G的x坐标
        mpz_init_set_str(g_y, "483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8", 16); // 基点G的y坐标
        mpz_init_set_str(n, "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141", 16); // 基点G的阶n
        mpz_init_set_str(h, "1", 10);        // 余因子h = 1
    }

    // 构造椭圆曲线对象
    EllipticCurve curve(a, b, p, g_x, g_y, n, h);

    {
        ZoneScopedN("ECDSA::CurveParams::get_secp256k1::Cleanup");
        // 清理临时变量
        mpz_clear(a); mpz_clear(b); mpz_clear(p);
        mpz_clear(g_x); mpz_clear(g_y);
        mpz_clear(n); mpz_clear(h);
    }

    return curve;
}

// prime256v1 曲线参数获取函数
EllipticCurve get_prime256v1(){
    ZoneScopedN("ECDSA::CurveParams::get_prime256v1"); // 标记曲线参数获取函数
    mpz_t a, b, p, g_x, g_y, n, h;

    {
        ZoneScopedN("ECDSA::CurveParams::get_prime256v1::Initialization");
        // prime256v1 (NIST P-256) 曲线参数
        mpz_init_set_str(a, "-3", 16); // 曲线参数 a
        mpz_init_set_str(b, "5AC635D8AA3A93E7B3EBBD55769886BC651D06B0CC53B0F63BCE3C3E27D2604B", 16); // 曲线参数 b
        mpz_init_set_str(p, "FFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFF", 16); // 素数域p
        mpz_init_set_str(g_x, "6B17D1F2E12C4247F8BCE6E563A440F277037D812DEB33A0F4A13945D898C296", 16); // 基点G的x坐标
        mpz_init_set_str(g_y, "4FE342E2FE1A7F9B8EE7EB4A7C0F9E162BCE33576B315ECECBB6406837BF51F5", 16); // 基点G的y坐标
        mpz_init_set_str(n, "FFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551", 16); // 基点G的阶n
        mpz_init_set_str(h, "1", 10); // 余因子h = 1
    }

    EllipticCurve curve(a, b, p, g_x, g_y, n, h);

    {
        ZoneScopedN("ECDSA::CurveParams::get_prime256v1::Cleanup");
        mpz_clear(a); mpz_clear(b); mpz_clear(p);
        mpz_clear(g_x); mpz_clear(g_y);
        mpz_clear(n); mpz_clear(h);
    }

    return curve;
}

// secp384r1 曲线参数获取函数
EllipticCurve get_secp384r1(){
    ZoneScopedN("ECDSA::CurveParams::get_secp384r1"); // 标记曲线参数获取函数
    mpz_t a, b, p, g_x, g_y, n, h;

    {
        ZoneScopedN("ECDSA::CurveParams::get_secp384r1::Initialization");
        // secp384r1 (NIST P-384) 曲线参数
        mpz_init_set_str(a, "fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff0000000000000000fffffffc", 16); // 曲线参数 a
        mpz_init_set_str(b, "b3312fa7e23ee7e4988e056be3f82d19181d9c6efe8141120314088f5013875ac656398d8a2ed19d2a85c8edd3ec2aef", 16); // 曲线参数 b
        mpz_init_set_str(p, "fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff0000000000000000ffffffff", 16); // 素数域p
        mpz_init_set_str(g_x, "aa87ca22be8b05378eb1c71ef320ad746e1d3b628ba79b9859f741e082542a385502f25dbf55296c3a545e3872760ab7", 16); // 基点G的x坐标
        mpz_init_set_str(g_y, "3617de4a96262c6f5d9e98bf9292dc29f8f41dbd289a147ce9da3113b5f0b8c00a60b1ce1d7e819d7a431d7c90ea0e5f", 16); // 基点G的y坐标
        mpz_init_set_str(n, "ffffffffffffffffffffffffffffffffffffffffffffffffc7634d81f4372ddf581a0db248b0a77aecec196accc52973", 16); // 基点G的阶n
        mpz_init_set_str(h, "1", 10); // 余因子h = 1
    }

    EllipticCurve curve(a, b, p, g_x, g_y, n, h);

    {
        ZoneScopedN("ECDSA::CurveParams::get_secp384r1::Cleanup");
        mpz_clear(a); mpz_clear(b); mpz_clear(p);
        mpz_clear(g_x); mpz_clear(g_y);
        mpz_clear(n); mpz_clear(h);
    }

    return curve;
}


template<typename curve_name>
void ECDSA<curve_name>::generatepublicparameters(ECDSA_publicparameter& publicParameters) {
    ZoneScopedN("ECDSA::generatepublicparameters"); // 标记整个函数

    // 使用 if constexpr 根据模板参数选择曲线
    if constexpr (std::is_same_v<curve_name, secp256k1_tag>) {
        ZoneScopedN("ECDSA::generatepublicparameters::Load secp256k1");
        // 选择 secp256k1 曲线
        publicParameters = ECDSA_publicparameter(get_secp256k1()); // 内部已插桩
        publicParameters.is_valid = true;
        std::cout << "Using secp256k1 curve parameters." << std::endl;
    } else if constexpr (std::is_same_v<curve_name, prime256v1_tag>) {
        ZoneScopedN("ECDSA::generatepublicparameters::Load prime256v1");
        // 选择 prime256v1 曲线
        publicParameters = ECDSA_publicparameter(get_prime256v1()); // 内部已插桩
        publicParameters.is_valid = true;
        std::cout << "Using prime256v1 curve parameters." << std::endl;
    } else if constexpr (std::is_same_v<curve_name, secp384r1_tag>) {
        ZoneScopedN("ECDSA::generatepublicparameters::Load secp384r1");
        // 选择 secp384r1 曲线
        publicParameters = ECDSA_publicparameter(get_secp384r1()); // 内部已插桩
        publicParameters.is_valid = true;
        std::cout << "Using secp384r1 curve parameters." << std::endl;
    } else {
        ZoneScopedN("ECDSA::generatepublicparameters::Unknown Curve Error");
        // 处理未知的曲线类型
        std::cerr << "Error: Unknown curve type provided to ECDSA template." << std::endl;
        publicParameters.is_valid = false;
        // 实际应用中可能需要抛出异常
    }
}

template<typename curve_name>
void ECDSA<curve_name>::generateKeyPair(
    ECDSA_privatekey& privateKey,
    ECDSA_publickey& publicKey,
    ECDSA_publicparameter& publicParameters) 
{
    ZoneScopedN("ECDSA::generateKeyPair"); // 标记整个函数

    {
        ZoneScopedN("ECDSA::generateKeyPair::Parameter Validation");
        if(!publicParameters.is_valid){
            generatepublicparameters(publicParameters); // 内部已插桩
        }
    }

    // 获取椭圆曲线参数
    const EllipticCurve& curve = publicParameters.ellipticcurve;
    
    // ------------------------------
    // 1. 生成私钥d：随机数 ∈ [1, n-1]
    // ------------------------------
    mpz_t n,d;
    gmp_randstate_t state;
    {
        ZoneScopedN("ECDSA::generateKeyPair::Private Key Generation");
        mpz_inits(n,d, nullptr);

        mpz_set(n, curve.n);  // 基点G的阶

        gmp_randinit_default(state);
        gmp_randseed_ui(state, static_cast<unsigned long>(time(nullptr)));
        
        do {
            mpz_urandomm(d, state, n); // 生成 [0, n-1)
        } while (mpz_cmp_ui(d, 0) == 0);         // 确保 d ≠ 0
    }
#ifdef DEBUG_PRINTS_LEVEL_3
    std::cout << "[ECDSA generateKeyPair][debug] Private key d: "; mpz_out_str(stdout, 10, d); std::cout << std::endl;
#endif

    {
        ZoneScopedN("ECDSA::generateKeyPair::Store Private Key and Cleanup");
        privateKey = ECDSA_privatekey(d);
        mpz_clears(n,d, nullptr);
        gmp_randclear(state);
    }

    // ------------------------------
    // 2. 计算公钥Q = d * G
    // ------------------------------
    EllipticCurve_point Q;
    {
        ZoneScopedN("ECDSA::generateKeyPair::Calculate Public Key Q");
        Q = scalar_mult(privateKey.d, curve.g, curve); // 内部已插桩
    }

    // 存储公钥
    {
        ZoneScopedN("ECDSA::generateKeyPair::Store Public Key");
        publicKey = ECDSA_publickey(Q);
    }
}

template<typename curve_name>
ECDSA_signature ECDSA<curve_name>::sign(
    const PlainMessage& message, 
    const ECDSA_privatekey& privatekey, 
    const ECDSA_publicparameter& publicparameter){

    ZoneScopedN("ECDSA::sign"); // 标记整个函数

    // 验证公共参数有效性
    {
        ZoneScopedN("ECDSA::sign::Parameter Validation");
        if (!publicparameter.is_valid) {
            throw std::invalid_argument("Invalid ECDSA public parameters");
        }
    }

    // 获取曲线参数
    const EllipticCurve& curve = publicparameter.ellipticcurve;
    const GMPInt& n = curve.n;
    const GMPInt& d = privatekey.d;

    // --------------------------
    // 1: 计算消息哈希z，转换为整数
    // --------------------------
    mpz_t z_raw;
    GMPInt z;
    {
        ZoneScopedN("ECDSA::sign::Hash Message to Integer");
        mpz_init(z_raw);
        // 转换哈希消息
        std::string hex_str;
        hash_value_to_hex_string(message.get_hash_data(), hex_str);
        int ret = mpz_set_str(z_raw, hex_str.c_str(), 16);
        if (ret != 0) {
            mpz_clear(z_raw);
            throw std::runtime_error("Invalid hash format");
        }
        z = GMPInt(z_raw); // 将z_raw转换为GMPInt
    }

    // ---------------
    // 2: 生成临时密钥k
    // ---------------
    GMPInt k;
    GMPInt r, s;
    gmp_randstate_t state;
    {
        ZoneScopedN("ECDSA::sign::Random State Initialization");
        gmp_randinit_default(state);
        gmp_randseed_ui(state, static_cast<unsigned long>(time(nullptr)));
    }
    
    {
        ZoneScopedN("ECDSA::sign::Signature Loop");
        do {
            // -------------------------
            // 3: 生成随机数k ∈ [1, n-1]
            // -------------------------
            {
                ZoneScopedN("ECDSA::sign::Generate K");
                do {
                    mpz_urandomm(k, state, n);
                } while (k == 0);
            }
            
            // -------------------------------------
            // 4: 计算椭圆曲线点kG(x_1, y_1) = k * G
            // -------------------------------------
            EllipticCurve_point kG;
            {
                ZoneScopedN("ECDSA::sign::Calculate kG");
                kG = scalar_mult(k, curve.g, curve); // 内部已插桩
            }

            // 验证点有效性
            if (kG.is_infinity) continue;

            // --------------------
            // 5: 计算r = x1 mod n
            // --------------------
            {
                ZoneScopedN("ECDSA::sign::Calculate R");
                r = kG.x % n;
            }
            if (r == 0) continue;

            // -----------------------------------
            // 6: 计算s = (z + r * d) * k⁻¹ mod n
            // -----------------------------------
            GMPInt k_inv;
            {
                ZoneScopedN("ECDSA::sign::Calculate K_inv");
                inverse_mod(k_inv, k, n); // 内部已插桩
            }
            {
                ZoneScopedN("ECDSA::sign::Calculate S");
                s = ((z + r * d) * k_inv) % n;
            }

        } while (s == 0);
    }

    // 清理随机状态
    {
        ZoneScopedN("ECDSA::sign::Cleanup Resources");
        gmp_randclear(state);
        mpz_clear(z_raw); // 清理原始的z_raw
    }

    return ECDSA_signature(r, s);
}

template<typename curve_name>
bool ECDSA<curve_name>::verify(
    const PlainMessage& message, 
    const ECDSA_publickey &publickey,
    const ECDSA_publicparameter &publicparameter, 
    const ECDSA_signature &signature){

    ZoneScopedN("ECDSA::verify"); // 标记整个函数

    // 验证公共参数有效性
    {
        ZoneScopedN("ECDSA::verify::Parameter Validation");
        if (!publicparameter.is_valid) {
            throw std::invalid_argument("Invalid ECDSA public parameters");
        }
    }

    mpz_t temp;
    mpz_init(temp);

    // 获取曲线参数
    const EllipticCurve& curve = publicparameter.ellipticcurve;
    const GMPInt& n = curve.n;

    // --------------
    // 1: 检查签名范围
    // --------------
    {
        ZoneScopedN("ECDSA::verify::Signature Range Check");
        mpz_set(temp,curve.n);
        if (signature.r <= 0 || mpz_cmp(signature.r, temp) >= 0 || // r < n
            signature.s <= 0 || mpz_cmp(signature.s, temp) >= 0) { // s < n
            mpz_clear(temp);
            return false;
        }
        mpz_clear(temp); // temp 仅用于检查范围，可以提前清理
    }


    // -------------------------------
    // 2: 计算消息哈希并将哈希转换为数值z
    // -------------------------------
    HASH_VALUE hash_val; // 避免与Z混淆
    GMPInt z;
    {
        ZoneScopedN("ECDSA::verify::Hash Message to Integer");
        hash_val = message.get_hash_data();
        std::string hex_str;
        hash_value_to_hex_string(hash_val, hex_str);
        mpz_t z_raw;
        mpz_init(z_raw);
        mpz_set_str(z_raw, hex_str.c_str(), 16);
        z = GMPInt(z_raw);
        mpz_clear(z_raw);
    }

    // -----------------------
    // 3: 计算逆元w = s⁻¹ mod n
    // -----------------------
    GMPInt w;
    {
        ZoneScopedN("ECDSA::verify::Calculate W (s_inverse)");
        inverse_mod(w, signature.s, n); // 内部已插桩
    }

    // -----------------
    // 4: 计算中间值u1/u2
    // -----------------
    GMPInt u1, u2;
    {
        ZoneScopedN("ECDSA::verify::Calculate U1 and U2");
        u1 = (z * w) % n;    // u1 = (z * s⁻¹) mod n
        u2 = (signature.r * w) % n;    // u2 = (r * s⁻¹) mod n
    }

    // -----------------------------------
    // 5: 计算椭圆曲线点P = u1 * G + u2 * Q
    // -----------------------------------
    EllipticCurve_point P1, P2, P;
    {
        ZoneScopedN("ECDSA::verify::Calculate P1 and P2");
        P1 = scalar_mult(u1, curve.g, curve); // 内部已插桩
        P2 = scalar_mult(u2, publickey.Q, curve); // 内部已插桩
    }
    {
        ZoneScopedN("ECDSA::verify::Add P1 and P2");
        P = point_add(P1, P2, curve); // 内部已插桩
    }

    // --------------
    // 6: 验证点有效性
    // --------------
    {
        ZoneScopedN("ECDSA::verify::Check P Validity");
        if (P.is_infinity) return false;
    }

    // -----------------
    // 7: 比较点的x坐标值
    // -----------------
    GMPInt x_P_mod_n;
    {
        ZoneScopedN("ECDSA::verify::Compare X-coordinate");
        x_P_mod_n = P.x % n;
        return (x_P_mod_n == signature.r);
    }
}


template class ECDSA<secp256k1_tag>;
template class ECDSA<prime256v1_tag>;
template class ECDSA<secp384r1_tag>;